#This script is working on PID control,slidebar mass should set to 0 for fix on GND
import pybullet as p  
#import pybullet_data  
import numpy as np  
import time
# 初始化PyBullet  
p.connect(p.GUI)  
p.setGravity(0,0,-9.8)
#p.setAdditionalSearchPath(pybullet_data.getData_path())  
p.setAdditionalSearchPath("D:\\f\\Open Source\\Prj_Robdog\\Simulation_Env\\bullet3\\data")
# 加载环境  
plane_id = p.loadURDF("plane.urdf") 
#cartpole_id = p.loadURDF("cartpole.urdf", [0, 0, 0.5]) # use system file 
#cartpole_id = p.loadURDF(".\\cartpole.urdf", [0, 0, 0.5])  # use local file, udpate pole mass
#cartpole_id = p.loadURDF(".\\cartpole_org.urdf", [0, 0, 0.5])   
cartpole_id = p.loadURDF(".\\my_cartpole.urdf", [0, 0, 0.5])  # use local file, udpate pole mass  
 

# Adjust friction for the cart and the plane  
p.changeDynamics(cartpole_id, -1, lateralFriction=0.05)  # Cart base with lower lateral friction 
p.changeDynamics(cartpole_id, 0, lateralFriction=0.05)  # Cart base with lower lateral friction  
p.changeDynamics(cartpole_id, 1, lateralFriction=0.05)  # pole with lower lateral friction 
p.changeDynamics(plane_id, -1, lateralFriction=0.05)     # Plane with lower lateral friction  
p.changeDynamics(plane_id, 0, lateralFriction=0.05)     # Plane with lower lateral friction 
# Set damping for the cart and pole (example values)  
p.changeDynamics(cartpole_id, -1, linearDamping=0.0, angularDamping=0.0)  # For the cart  
p.changeDynamics(cartpole_id, 0, linearDamping=0.0, angularDamping=0.0)  # For the cart  
p.changeDynamics(cartpole_id, 1, linearDamping=0.0, angularDamping=0.0)  # For the pole

# Alternatively, if the pole is a child of the cart:  
# Set joint state  
initial_angle = 0.5  # or any appropriate angle  
p.resetJointState(cartpole_id, 1, initial_angle)  


 

# MPC参数  
N = 10  # 预测步长  
#K = 1000  # 控制增益  

def get_state():  
    pos, _ = p.getBasePositionAndOrientation(cartpole_id)  
    angle = p.getJointState(cartpole_id, 1)[0]  # 杆的角度  
    return np.array([pos[0], angle])  



# Initialize variables for PID controller 
# PID parameters  
Kp = 1000  # Proportional gain  
Ki = 10  # Integral gain  
Kd = 20  # Derivative gain, senstive to dt   
integral = 0.0  
previous_error = 0.0  
# Set clamp limits for integral term  
integral_max = 0.05  # Maximum value for the integral  
integral_min = -0.05 # Minimum value for the integral  

def pid_control(state, target_angle, dt):  
    global integral, previous_error ,integral_max, integral_min
    
    x, theta = state  # Unpack the state 
    angle_error = -target_angle +  theta  # Calculate the error 
    #add gain 
    if abs(angle_error) >= 0.05:  
        angle_error = angle_error* 50
       
    
    # Proportional term  
    P_term = Kp * angle_error  
    
    # Integral term  
    integral += angle_error * dt 
    #avoid integral windup 
    if integral > integral_max:  
        integral = integral_max  
    elif integral < integral_min:  
        integral = integral_min 

    
    I_term = Ki * integral  
    
    # Derivative term  
    derivative = (angle_error - previous_error) / dt  
    D_term = Kd * derivative  
    
    # Calculate control input  
    control_input = P_term + I_term + D_term  
    
    # Clip the control input to prevent excessive commands, large mass need large force
    return np.clip(control_input, -4000, 4000)  




def print_dynamics_info(object_id, description):  
    for link_index in range(p.getNumJoints(object_id)):  
        dynamics_info = p.getDynamicsInfo(object_id, link_index)  
        print(f"{description} - Link:{link_index}, Mass: {dynamics_info[0]}, Friction: {dynamics_info[1]}")  

print_dynamics_info(cartpole_id, "Cartpole")  
print_dynamics_info(plane_id, "Plane")

# Print out joint damping information  
for i in range(p.getNumJoints(cartpole_id)):  
    joint_info = p.getJointInfo(cartpole_id, i)  
    print(f"Joint {i}: Damping {joint_info[6]}")  # Joint damping coefficient




# 设置仿真参数  
num_steps = 100 #need adjust freely  
# Simulation loop
try:
    while True:
        
        update_state = get_state()  
        
        # 计算控制输入 , here time should be sleep time * num_steps(interval of 2 states), 
        #dt change need udpate PID D , senstive to dt
        #1500N only keep pole angle not change , only 2000N make difference !!!
        control_input = pid_control(update_state,0.0,0.1)
        #control_input = 2000    # pole set to 0.5kg, 250N to make cart move, 10kg, 400N to make cart move
        print("control input : ", control_input)

        # 应用控制输入 
        #for step in range(num_steps):  
        
        p.setJointMotorControl2(cartpole_id, 0, p.TORQUE_CONTROL, force=control_input) 
        # Apply control input as an external force to the cart  
        # p.applyExternalForce(objectUniqueId=cartpole_id,   
        # linkIndex=0,  # Base (cart)  
        # forceObj=[control_input, 0, 0],   
        # posObj=p.getBasePositionAndOrientation(cartpole_id)[0],   
        # flags=p.WORLD_FRAME)  
        for step in range(num_steps):
            # 前进一步, pole fall takes time,  
            p.stepSimulation()  
            #time.sleep(0.01)
            
        update_state = get_state()             
        # 打印状态  
        #if step % 100 == 0:  
        print(f" Position: {update_state[0]:.2f}, Angle: {update_state[1]:.2f}") 
               
        
        
         
             
             
                
except KeyboardInterrupt:
    # Disconnect from the physics engine
    p.disconnect()


###Problem , pole didn't fall naturally at init angle, only adding force on cart can change
## Tips: 
# PID control, set max force to 2000 instead of 1000, pole become stable.(大力出奇迹)
## Set intergral max and min to avoid integral windup, pole become more stable.
# Pole mass set to 10,control foce need 2000N.  if it's 0.5, not stable
## cart mass also matters, or PID need adjustment, so it's possible to tune PID in simulator first
## slidebar mass set to 0 to make it fix on GND
## Need define pole rotion axis, or pole will not fall
## pole fall speed -> large mass & inertial fall faster, small sleep number fall fast 
## need set slidebar mass to 0, or it will move 
##1500N only keep pole angle not change , only 2000N make difference , so set force at least 2000N!!!
## I become PID Master now:
## Pole become very stable in short time,
# when angle error is small, multiple gain to amplify angle error, so do not a large PID gain,
# # which make pole stable when err is small

###Push into stack
# P control is ok, but need > 250 N to make cart move , is it a problem? may pybullet char 
#cofnirmed it is happened after adding sliderbar
#it's a proble, pole not fall at all. -> didn't set rotation axis  
# PID control, set big P to make force > 2000 , since dt change from 0.01 to 0.1, so need to change PID
##MPC controller output same value -> normal as pole angle not change, states not changed, output same.
#need make pole fall quickly first to change status...